package com.company.algo.SearchAlgorithm.BFS;

import com.company.algo.pojo.TreeNode;

import java.util.*;

/**
 * 863. 二叉树中所有距离为 K 的结点
 */
public class DistanceK {
    //key-顶点，value-邻居
    HashMap<Integer,List<Integer>> graph = new HashMap<>();
    //这道题就是先把二叉树转化图,再用图的bfs,求得解
    public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
        buildGraph(root);
        Queue<Integer> queue = new LinkedList<>();
        queue.offer(target.val);
        boolean[] visited = new boolean[500];
        visited[target.val]=true;
        List<Integer> res= new ArrayList<>();
        while(k>=0){
            res = new ArrayList<>();
            for (int i = queue.size(); i > 0 ; i--) {
                int cur = queue.poll();
                res.add(cur);
                for(int neighbor: graph.get(cur)){
                    if(!visited[neighbor]){
                        visited[neighbor] = true;
                        queue.offer(neighbor);
                    }
                }
            }
            k -= 1;
        }
        return res;
    }

    public void buildGraph(TreeNode root){
        graph.putIfAbsent(root.val, new ArrayList<>());
        if(root.left!=null){
            graph.get(root.val).add(root.left.val);
            graph.putIfAbsent(root.left.val, new ArrayList<>());
            graph.get(root.left.val).add(root.val);
            buildGraph(root.left);
        }
        if(root.right!=null){
            graph.get(root.val).add(root.right.val);
            graph.putIfAbsent(root.right.val, new ArrayList<>());
            graph.get(root.right.val).add(root.val);
            buildGraph(root.right);
        }
    }

    public static void main(String[] args) {
        TreeNode node3 = new TreeNode(3);
        TreeNode node4 = new TreeNode(4);
        TreeNode node2 = new TreeNode(2,null,node4);
        TreeNode node1 = new TreeNode(1,node3,null);
        TreeNode root = new TreeNode(0,node1,node2);
        DistanceK Main = new DistanceK();
        System.out.println(Main.distanceK(root, node2, 2));

    }
}
